// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

namespace kumo::pollux::parquet {
    class StringDecoder {
    public:
        StringDecoder(const char *start, const char *end, int fixedLength = -1)
            : bufferStart_(start),
              lastSafeWord_(end - simd::kPadding),
              fixedLength_(fixedLength) {
        }

        void skip(uint64_t numValues) {
            skip<false>(numValues, 0, nullptr);
        }

        template<bool hasNulls>
        inline void skip(int32_t numValues, int32_t current, const uint64_t *nulls) {
            if (hasNulls) {
                numValues = bits::countNonNulls(nulls, current, current + numValues);
            }
            for (auto i = 0; i < numValues; ++i) {
                bufferStart_ += lengthAt(bufferStart_) + sizeof(int32_t);
            }
        }

        template<bool hasNulls, typename Visitor>
        void readWithVisitor(const uint64_t *nulls, Visitor visitor) {
            int32_t current = visitor.start();
            int32_t numValues = 0;
            skip<hasNulls>(current, 0, nulls);
            int32_t toSkip;
            bool atEnd = false;
            const bool allowNulls = hasNulls && visitor.allowNulls();
            for (;;) {
                if (hasNulls && allowNulls && bits::isBitNull(nulls, current)) {
                    toSkip = visitor.processNull(atEnd);
                } else {
                    if (hasNulls && !allowNulls) {
                        toSkip = visitor.checkAndSkipNulls(nulls, current, atEnd);
                        if (!Visitor::dense) {
                            skip<false>(toSkip, current, nullptr);
                        }
                        if (atEnd) {
                            if constexpr (Visitor::kHasHook) {
                                visitor.setNumValues(
                                    Visitor::kHasFilter ? numValues : visitor.numRows());
                            }
                            return;
                        }
                    }

                    // We are at a non-null value on a row to visit.
                    toSkip = visitor.process(
                        fixedLength_ > 0 ? readFixedString() : readString(), atEnd);
                }
                ++current;
                ++numValues;
                if (toSkip) {
                    skip<hasNulls>(toSkip, current, nulls);
                    current += toSkip;
                }
                if (atEnd) {
                    if constexpr (Visitor::kHasHook) {
                        visitor.setNumValues(
                            Visitor::kHasFilter ? numValues : visitor.numRows());
                    }
                    return;
                }
            }
        }

    private:
        int32_t lengthAt(const char *buffer) {
            return *reinterpret_cast<const int32_t *>(buffer);
        }

        melon::StringPiece readString() {
            auto length = lengthAt(bufferStart_);
            bufferStart_ += length + sizeof(int32_t);
            return melon::StringPiece(bufferStart_ - length, length);
        }

        melon::StringPiece readFixedString() {
            bufferStart_ += fixedLength_;
            return melon::StringPiece(bufferStart_ - fixedLength_, fixedLength_);
        }

        const char *bufferStart_;
        const char *const lastSafeWord_;
        const int fixedLength_;
    };
} // namespace kumo::pollux::parquet
